# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.

# SPDX-License-Identifier: Apache-2.0

import pytest
import ttnn

from tests.nightly.t3000.ccl.test_strided_all_gather_minimal_matmul_async import (
    run_strided_all_gather_minimal_matmul_impl,
)
from models.common.utility_functions import skip_for_blackhole


# tiles_per_chunk needs to be divisible by num_workers_per_link
# mm_cores_y is the number of in0 first col cores
# mm_block_h and mm_block_w is the mm_block of a single mm_core_y
# so the result of one chunk transfer will be mm_cores_y * mm_block_h * mm_block_w, which will be tiles_per_chunk.  tiles_per_chunk % num_workers_per_link must equal 0
@skip_for_blackhole("Requires wormhole_b0 to run")
@pytest.mark.parametrize("mesh_device", [(8, 4)], indirect=True)
@pytest.mark.parametrize(
    "M, K, N, dim, other_dim, num_links, num_workers_per_link, layout, ag_input_dtype, mm_block_m, mm_block_k, mm_block_n, subblock_h, subblock_w, mm_core_grid, shard_weights, ag_offset",
    [
        # (
        #     512,
        #     256,
        #     512,
        #     3,
        #     2,
        #     1,
        #     1,
        #     ttnn.TILE_LAYOUT,
        #     ttnn.bfloat16,
        #     32,
        #     32,
        #     32,
        #     1,
        #     1,
        #     ttnn.CoreCoord(2, 2),
        #     False,
        #     (0, 6),
        # ),
        # (
        #     512,
        #     512,
        #     512,
        #     3,
        #     2,
        #     2,
        #     1,
        #     ttnn.TILE_LAYOUT,
        #     ttnn.bfloat16,
        #     32,
        #     64,
        #     32,
        #     1,
        #     1,
        #     ttnn.CoreCoord(2, 2),
        #     False,
        #     (0, 6),
        # ),
        # (
        #     4096,
        #     4096,
        #     4096,
        #     3,
        #     2,
        #     1,
        #     1,
        #     ttnn.TILE_LAYOUT,
        #     ttnn.bfloat16,
        #     32,
        #     32,
        #     32,
        #     1,
        #     1,
        #     ttnn.CoreCoord(2, 2),
        #     False,
        #     (0, 6),
        # ),
        # (
        #     4096,
        #     4096,
        #     4096,
        #     3,
        #     2,
        #     4,
        #     1,
        #     ttnn.TILE_LAYOUT,
        #     ttnn.bfloat16,
        #     64,
        #     64,
        #     32,
        #     1,
        #     1,
        #     ttnn.CoreCoord(2, 2),
        #     False,
        #     (0, 6),
        # ),
        # (
        #     4096,
        #     4096,
        #     4096,
        #     3,
        #     2,
        #     1,
        #     1,
        #     ttnn.TILE_LAYOUT,
        #     ttnn.bfloat16,
        #     32,
        #     32,
        #     32,
        #     1,
        #     1,
        #     ttnn.CoreCoord(4, 4),
        #     False,
        #     (0, 6),
        # ),
        # (
        #     4096,
        #     4096,
        #     4096,
        #     3,
        #     2,
        #     4,
        #     1,
        #     ttnn.TILE_LAYOUT,
        #     ttnn.bfloat16,
        #     64,
        #     64,
        #     32,
        #     1,
        #     1,
        #     ttnn.CoreCoord(4, 4),
        #     False,
        #     (0, 6),
        # ),
        # (
        #     4096,
        #     4096,
        #     4096,
        #     3,
        #     2,
        #     1,
        #     2,
        #     ttnn.TILE_LAYOUT,
        #     ttnn.bfloat16,
        #     256,
        #     256,
        #     256,
        #     2,
        #     2,
        #     ttnn.CoreCoord(2, 2),
        #     False,
        #     (0, 5),
        # ),
        # (
        #     4096,
        #     4096,
        #     4096,
        #     3,
        #     2,
        #     4,
        #     2,
        #     ttnn.TILE_LAYOUT,
        #     ttnn.bfloat16,
        #     256,
        #     256,
        #     256,
        #     2,
        #     2,
        #     ttnn.CoreCoord(4, 4),
        #     False,
        #     (0, 5),
        # ),
        (
            75776,
            5120,
            1280,
            3,
            2,
            4,
            2,
            ttnn.TILE_LAYOUT,
            ttnn.bfloat16,
            256,
            256,
            256,
            2,
            2,
            ttnn.CoreCoord(8, 6),
            False,
            (0, 6),
        ),
        # (
        #     75776,
        #     5120,
        #     3456,
        #     3,
        #     2,
        #     4,
        #     2,
        #     ttnn.TILE_LAYOUT,
        #     ttnn.bfloat16,
        #     256,
        #     256,
        #     256,
        #     2,
        #     2,
        #     ttnn.CoreCoord(8, 6),
        #     False,
        #     (0, 6),
        # ),
    ],
    ids=[
        # "base1link",  # 1 forward pass through K
        # "base2link",  # 1 forward pass through K
        # "4k4k4k1link",
        # "4k4k4k4link",
        # "4x4mmcores1link",
        # "4x4mmcores4link",
        # "fulltest1link",
        # "fulltest4link",
        "wan1",
        # "wan2",
    ],
)
@pytest.mark.parametrize(
    "mem_config_input, mem_config_ag, mem_config_mm",
    [
        (
            ttnn.MemoryConfig(ttnn.TensorMemoryLayout.INTERLEAVED, ttnn.BufferType.DRAM),
            ttnn.MemoryConfig(ttnn.TensorMemoryLayout.INTERLEAVED, ttnn.BufferType.DRAM),
            ttnn.MemoryConfig(ttnn.TensorMemoryLayout.INTERLEAVED, ttnn.BufferType.DRAM),
        )
    ],
)
@pytest.mark.parametrize(
    "enable_trace,num_iters",
    [
        (False, 1),
    ],
    ids=["check"],
)
@pytest.mark.parametrize(
    "use_non_fused",
    [
        False,
    ],
    ids=["fused"],
)
@pytest.mark.parametrize(
    "read_local_slice_from_input",
    [
        True,
    ],
    ids=["read_local"],
)
@pytest.mark.parametrize(
    "device_params, all_gather_topology",
    [
        ({"fabric_config": ttnn.FabricConfig.FABRIC_1D_RING, "trace_region_size": 90112}, ttnn.Topology.Ring),
    ],
    indirect=["device_params"],
    ids=["fabric_ring"],
)
def test_strided_all_gather_minimal_matmul_async(
    mesh_device,
    M,
    K,
    N,
    dim,
    other_dim,
    num_links,
    ag_input_dtype,
    layout,
    mem_config_input,
    mem_config_ag,
    mem_config_mm,
    enable_trace,
    all_gather_topology,
    num_iters,
    num_workers_per_link,
    mm_block_m,
    mm_block_k,
    mm_block_n,
    subblock_h,
    subblock_w,
    mm_core_grid,
    use_non_fused,
    shard_weights,
    ag_offset,
    read_local_slice_from_input,
):
    run_strided_all_gather_minimal_matmul_impl(
        mesh_device,
        mesh_device.get_num_devices(),
        M,
        K,
        N,
        dim,
        other_dim,
        num_links,
        ag_input_dtype,
        layout,
        mem_config_input,
        mem_config_ag,
        mem_config_mm,
        all_gather_topology=all_gather_topology,
        enable_trace=enable_trace,
        num_iters=num_iters,
        num_workers_per_link=num_workers_per_link,
        mm_block_m=mm_block_m,
        mm_block_k=mm_block_k,
        mm_block_n=mm_block_n,
        subblock_h=subblock_h,
        subblock_w=subblock_w,
        mm_core_grid=mm_core_grid,
        use_non_fused=use_non_fused,
        shard_weights=shard_weights,
        ag_core_grid_offset=ag_offset,
        read_local_slice_from_input=read_local_slice_from_input,
    )
